import pandas as pd
import matplotlib.pyplot as plt
import numpy as np



class StockMetricsCalculator:
    def __init__(self, ohlcv_data):
        """
        Args:
            ohlcv_data: pandas DataFrame 必须包含以下列:
                ['open', 'high', 'low', 'close', 'vol']
        """
        self.data = ohlcv_data

    def _validate_data(self, data, required):
        # required = ['open', 'high', 'low', 'close', 'vol']
        if not all(col in data.columns for col in required):
            raise ValueError(f"缺少必须的列，需要包含: {required}")

    def SMA(self, window, column='close'):
        """计算简单移动平均线"""
        self._validate_data(self.data, [column])
        data = self.data[column]
        return data.rolling(window).mean()

    def EMA(self, window, column='close'):
        """计算指数移动平均线"""
        self._validate_data(self.data, [column])
        data = self.data[column]
        return data.ewm(span=window, adjust=False).mean()

    def RSI(self, period=14, column='close'):
        """计算相对强弱指数"""
        self._validate_data(self.data, [column])
        data = self.data[column]
        # 计算价格变化
        Change = data.diff()

        # 计算上涨和下跌
        Gain = Change.apply(lambda x: x if x > 0 else 0)
        Loss = Change.apply(lambda x: -x if x < 0 else 0)

        # 计算平均上涨和平均下跌
        Avg_Gain = Gain.rolling(window=period, min_periods=1).mean()
        Avg_Loss = Loss.rolling(window=period, min_periods=1).mean()

        # 计算相对强度RS
        RS = Avg_Gain / Avg_Loss

        # 计算RSI返回
        return 100 - (100 / (1 + RS))

    def Bollinger_bands(self, window=20, num_std=2, column='close'):
        """计算布林带
        Returns:
            (middle, upper, lower) 元组 N日中轨线，上轨线，下轨线
        """
        self._validate_data(self.data, [column])
        sma = self.SMA(window, column)
        std = self.data[column].rolling(window).std()
        upper = sma + (std * num_std)
        lower = sma - (std * num_std)
        return sma, upper, lower

        def VWAP(self):
        self._validate_data(self.data, ['high', 'low', 'close', 'vol'])
        """计算成交量加权平均价
        将多笔交易的价格按各自的成交量加权而算出的平均价"""
        cum_volume = self.data['vol'].cumsum()
        cum_price_volume = (self.data[['high', 'low', 'close']].mean(axis=1) * self.data['vol']).cumsum()
        return cum_price_volume / cum_volume

    def Historical_volatility(self, window=252, column='close'):
        self._validate_data(self.data, [column])
        """计算历史波动率过去一段时间的变化"""
        log_returns = np.log(self.data[column] / self.data[column].shift(1))
        return log_returns.rolling(window).std() * np.sqrt(window)

    def Standard_deviation(self, window, column='close'):
        self._validate_data(self.data, [column])
        """计算指定窗口期的价格标准差"""
        return self.data[column].rolling(window).std()

    # 可选扩展指标
    def OBV(self):
        self._validate_data(self.data, ['close', 'vol'])
        """计算能量潮指标"""
        obv = (np.sign(self.data['close'].diff()) * self.data['vol']).cumsum()
        return obv.fillna(0)

    def ATR(self, window=14):
        """计算window长度的平均真实波幅"""
        self._validate_data(self.data, ['high', 'low', 'close'])
        high_low = self.data['high'] - self.data['low']
        high_close = np.abs(self.data['high'] - self.data['close'].shift())
        low_close = np.abs(self.data['low'] - self.data['close'].shift())
        true_range = pd.concat([high_low, high_close, low_close], axis=1).max(axis=1)
        return true_range.rolling(window).mean()

    def log_Return(self):
        """计算对数收益"""
        self._validate_data(self.data, ['close'])
        returns = pd.to_numeric(self.data['close'] / self.data['close'].shift(1),
                                errors='coerce')  # 将数据转化为数值类型 errors='coerce'：空为NaN
        log_returns = np.log(returns)
        return log_returns

    def DMI(self, period=14):
        """计算动量指数(DMI)
            返回 +DI, -DI, ADX, ADXR
        """
        self._validate_data(self.data, ['high', 'low', 'close'])
        # res = np.full(self.data.shape[0], np.nan)
        high = self.data['high']
        low = self.data['low']
        close = self.data['close']

        # 计算TR
        tr1 = high - low
        tr2 = abs(high - close.shift())
        tr3 = abs(low - close.shift())
        tr = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)

        # 计算+DM和-DM
        plus_dm = high - high.shift()
        minus_dm = low.shift() - low
        mask_plus = (plus_dm > minus_dm) & (plus_dm > 0)
        mask_minus = (minus_dm > plus_dm) & (minus_dm > 0)
        plus_dm = plus_dm * mask_plus
        minus_dm = minus_dm * mask_minus

        # 计算+DI和-DI
        smooth_tr = tr.ewm(span=period).mean()
        smooth_plus_dm = plus_dm.ewm(span=period).mean()
        smooth_minus_dm = minus_dm.ewm(span=period).mean()
        plus_di = (smooth_plus_dm / smooth_tr) * 100
        minus_di = (smooth_minus_dm / smooth_tr) * 100

        # 计算ADX
        diff_di = abs(plus_di - minus_di)
        sum_di = plus_di + minus_di
        dx = (diff_di / sum_di) * 100
        adx = dx.ewm(span=period).mean()

        # 计算ADXR
        adxr = ((adx.shift(period) + adx) / 2)

        # +DI, -DI, ADX, ADXR
        return plus_di, minus_dm, adx, adxr

    def TRIX(self, period=30, m=9):
        """计算三重指数平滑移动平均(TRIX) 中长线按指标
        Returns:
            trix 三重指数平滑移动平均(TRIX) N日的TRIX"""

        self._validate_data(self.data, ['close'])
        close = self.data['close']

        # 1. 计算N天的收盘价的指数平均AX
        ax = close.ewm(span=period, adjust=False).mean()

        # 2. 计算N天的AX的指数平均BX
        bx = ax.ewm(span=period, adjust=False).mean()

        # 3. 计算N天的BX的指数平均TRIX
        trix = bx.ewm(span=period, adjust=False).mean()

        # 计算TRIX的百分比变化率
        trix = trix.pct_change() * 100

        # 4. 计算TRIX的m日移动平均TRMA
        trma = trix.rolling(window=m).mean()

        # TRIX, TRMA
        return trix, trma

    def DMA(self, short_window=10, long_window=50):
        """计算动态移动平均线(DMA)
        Returns:
            DMA: 长期窗口"""

        self._validate_data(self.data, ['close'])
        close = self.data['close']

        # 计算短期移动平均线
        short_sma = close.rolling(window=short_window).mean()

        # 计算长期移动平均线
        long_sma = close.rolling(window=long_window).mean()

        # 计算DMA
        dma = short_sma - long_sma

        # 返回DMA的Numpy数组形式，同时去除可能存在的NaN值
        return dma

    def W_R(self, period=14):
        """计算威威廉指标(Williams R)
        Returns:
            WR： 威廉指标
        """
        self._validate_data(self.data, ['high', 'low', 'close'])
        high_prices = self.data['high']
        low_prices = self.data['low']
        close_prices = self.data['close']

        highest_high = high_prices.rolling(window=period).max()
        lowest_low = low_prices.rolling(window=period).min()

        wr = ((highest_high - close_prices) / (highest_high - lowest_low)) * 100
        return wr

    def BIAS(self, period=10):
        """计算乖离率(BIAS)
        """
        self._validate_data(self.data, ['close'])
        # 计算N日简单移动平均
        ma_N = self.data['close'].rolling(window=period).mean()

        # 计算BIAS
        bias = ((self.data['close'] - ma_N) / ma_N) * 100

        # 将结果转换为NumPy数组并返回
        return bias

    def CR(self, period=10):
        """计算中间意愿指标(CR)"""
        self._validate_data(self.data, ['high', 'low'])
        high = self.data['high']
        low = self.data['low']

        # 近period日的最高价与最低价的平均值
        high_avg = high.rolling(window=period).mean()
        low_avg = low.rolling(window=period).mean()

        # 前一日的最高价与最低价
        recent_high = high.shift(1)
        recent_low = low.shift(1)

        # CR指标计算
        cr = (recent_high + recent_low) / 2
        cr = cr / low_avg * 100
        cr = cr / high_avg * low_avg
        return cr

    def VR(self, short_period=26, long_period=52):
        """计算成交量变异率(VR)"""

        self._validate_data(self.data, ['close', 'pre_close'])
        # 计算每日涨跌
        daily_change = self.data['close'] - self.data['pre_close']

        # 初始化累积成交量数组
        up_vol = np.zeros(len(self.data))
        down_vol = np.zeros(len(self.data))

        # 遍历数据，计算上涨和下跌日的成交量
        for i in range(1, len(self.data)):
            if daily_change.iloc[i] > 0 or (daily_change.iloc[i] == 0 and daily_change.iloc[i - 1] > 0):
                up_vol[i] = self.data['vol'].iloc[i]
            elif daily_change.iloc[i] < 0 or (daily_change.iloc[i] == 0 and daily_change.iloc[i - 1] < 0):
                down_vol[i] = self.data['vol'].iloc[i]

        # 计算短期和长期的平均成交量
        up_vol_short = pd.Series(up_vol).rolling(window=short_period).sum()
        down_vol_long = pd.Series(down_vol).rolling(window=long_period).sum()

        # 计算VR指标
        vr = up_vol_short / down_vol_long

        # 转换为NumPy数组并返回
        return vr

    def MIKE(self, period=10):
        """计算麦克指标(MIKE)"""
        self._validate_data(self.data, ['high', 'low', 'close'])
        high = self.data['high']
        low = self.data['low']
        close = self.data['close']

        # 计算起始价
        TYP = (high + low + close) / 3

        # period内的最高价和最低价
        maxx = high.rolling(window=period).max()
        minn = low.rolling(window=period).min()

        # 计算初级压力WR，中级压力MR，强力压力SR
        WR = TYP * 2 - minn
        MR = TYP + (maxx - minn)
        SR = 2 * maxx - minn

        # 计算初级支撑WS，中级支撑MS，强力支撑SS
        WS = TYP - (maxx - TYP)
        MS = TYP - (maxx - minn)
        SS = 2 * minn - maxx

        return np.stack([TYP, WR, MR, SR, WS, MS, SS], axis=1)

    def PVI(self, period=14):
        """计算为正成交量指标"""

        self._validate_data(self.data, ['close', 'vol'])
        # 计算收盘价和成交量的移动平均
        avg_close = self.data['close'].rolling(window=period).mean()
        avg_vol = self.data['vol'].rolling(window=period).mean()

        # 初始化PVI值
        pvi = np.zeros(len(self.data))
        pvi[0] = 100.0  # 第一次计算时，PV以100代替

        # 计算PVI
        for i in range(1, len(self.data)):
            # 当前收盘价与过去N日收盘价的平均值的差异
            close_diff = self.data['close'].iloc[i] - avg_close.iloc[i]

            # 当前成交量与过去N日成交量的平均值的差异
            vol_diff = self.data['vol'].iloc[i] - avg_vol.iloc[i]

            # 如果成交量差大于0，则使用公式计算新的PVI值
            if vol_diff > 0:
                pvi[i] = pvi[i - 1] + (close_diff / avg_close.iloc[i]) * pvi[i - 1]
            else:
                pvi[i] = pvi[i - 1]

        return pvi

    def ASI(self):
        """计算积累摆动指数"""
        self._validate_data(self.data, ['high', 'low', 'close', 'open'])
        df = self.data

        # 初始化ASI数组，长度与数据框行数相同，初始值为NaN
        asi = np.full(len(df), np.nan)

        # 如果数据框行数小于2，则直接返回全为NaN的数组
        if len(df) < 2:
            return asi

        # 计算ASI
        for i in range(1, len(df)):
            # 获取数据
            today_high = df['high'].iloc[i]
            today_low = df['low'].iloc[i]
            yesterday_close = df['close'].iloc[i - 1]
            yesterday_open = df['open'].iloc[i - 1]
            today_close = df['close'].iloc[i]
            today_open = df['open'].iloc[i]

            # 计算A、B、C、D、E、F、G
            A = abs(today_high - yesterday_close)
            B = abs(today_low - yesterday_close)
            C = abs(today_high - today_low)
            D = abs(yesterday_close - yesterday_open)
            E = today_close - yesterday_close
            F = today_close - today_open
            G = yesterday_close - yesterday_open

            # 计算X
            X = E + 0.5 * F + G

            # 比较A、B，取最大值
            maxx = max(A, B, C)
            if maxx == A:
                R = A + 0.5 * B + 0.25 * D
            elif maxx == B:
                R = B + 0.5 * A + 0.25 * D
            else:
                R = C + 0.25 * D

            # 计算K和L
            K = max(A, B)
            L = 13

            if (R == 0):
                R = 1e-3
            if (L == 0):
                L = 1e-3

            # 计算SI
            if R != 0 and L != 0:
                SI = 50 * (X / R) * (K / L)
            else:
                SI = np.nan

            # 计算ASI
            if i == 1:
                asi[i] = SI
            else:
                asi[i] = asi[i - 1] + SI

        return asi

    def KDJ(self, n=9):
        """计算KDJ指标"""
        self._validate_data(self.data, ['high', 'low', 'close'])
        df = self.data

        # 初始化K、D、J值的数组，使用NaN填充
        K = np.full_like(df['close'], np.nan)
        D = np.full_like(df['close'], np.nan)
        J = np.full_like(df['close'], np.nan)

        # 计算RSV
        df = self.data

        # 初始化K、D、J值的数组，使用NaN填充
        K = np.full_like(df['close'], np.nan)
        D = np.full_like(df['close'], np.nan)
        J = np.full_like(df['close'], np.nan)

        # 计算RSV
        for i in range(n - 1, len(df)):
            # 计算过去n天内的最高价和最低价
            high = df['high'].iloc[i - n + 1:i + 1].max()
            low = df['low'].iloc[i - n + 1:i + 1].min()

            # 计算当日的RSV
            if high == low:
                RSV = 50  # 避免除零错误，当最高价等于最低价时，RSV设为50
            else:
                RSV = (df['close'].iloc[i] - low) / (high - low) * 100

            # 计算K值和D值
            if i == n - 1:
                # 如果是第一天，没有前一日的K值和D值，用50代替
                K[i] = RSV
                D[i] = K[i]
            else:
                K[i] = 2 / 3 * K[i - 1] + 1 / 3 * RSV
                D[i] = 2 / 3 * D[i - 1] + 1 / 3 * K[i]

            # 计算J值
            J[i] = 3 * D[i] - 2 * K[i]

        return np.stack([K, D, J], axis=1)

    def ROC(self, n=1):
        """计算ROC变动率指标"""
        self._validate_data(self.data, ['close'])
        df = self.data

        # 初始化ROC数组，使用NaN填充
        ROC = np.full_like(df['close'], np.nan)

        # 计算ROC
        for i in range(n, len(df)):
            # 计算n天前的收盘价
            close_n_days_ago = df['close'].iloc[i - n]

            # 计算当日的ROC
            if close_n_days_ago == 0:
                close_n_days_ago = 1e-3

            ROC[i] = ((df['close'].iloc[i] - close_n_days_ago) / close_n_days_ago) * 100

        return ROC

    def MTM(self, n=3):
        """计算Momentum指标"""
        self._validate_data(self.data, ['close'])
        df = self.data

        MTM = np.full_like(df['close'], np.nan)

        for i in range(n, len(df)):
            MTM[i] = df['close'].iloc[i] - df['close'].iloc[i - n]

        return MTM

    def ARBR(self, n=26):
        """计算Adaptive Relative Strength Index (ARBR)"""
        self._validate_data(self.data, ['high', 'low', 'close', 'open'])
        df = self.data

        # 初始化AR和BR数组，使用NaN填充
        AR = np.full_like(df['close'], np.nan)
        BR = np.full_like(df['close'], np.nan)

        # 计算AR和BR
        df = self.data

        # 初始化AR和BR数组，使用NaN填充
        AR = np.full_like(df['close'], np.nan)
        BR = np.full_like(df['close'], np.nan)

        # 计算AR和BR
        for i in range(n - 1, len(df)):
            # 计算AR
            sum_HO = df['high'].iloc[i - n + 1:i + 1].sub(df['open'].iloc[i - n + 1:i + 1]).sum()
            sum_OL = df['open'].iloc[i - n + 1:i + 1].sub(df['low'].iloc[i - n + 1:i + 1]).sum()
            if sum_OL == 0:
                sum_OL = 1e-3
            # else:
            AR[i] = (sum_HO / sum_OL) * 100

            # 计算BR
            if i == 0:
                BR[i] = np.nan  # 第一天没有昨天的收盘价，无法计算BR
            else:
                sum_HPC = df['high'].iloc[i - n + 1:i + 1].sub(df['close'].iloc[i - 1]).sum()
                sum_PCL = df['close'].iloc[i - 1:i].sub(df['low'].iloc[i - n + 1:i + 1]).sum()
                if sum_PCL == 0:
                    sum_PCL = 1e-3
                # else:
                BR[i] = (sum_HPC / sum_PCL) * 100

        return np.stack([AR, BR], axis=1)

    def PSY(self, n=12):
        """计算心理线(PSY)指标"""
        self._validate_data(self.data, ['close'])
        df = self.data

        # 初始化PSY数组，使用NaN填充
        PSY = np.full_like(df['close'], np.nan)

        # 计算PSY
        for i in range(n, len(df)):
            # 计算N日内上涨天数
            up_days_count = 0
            for j in range(i - n, i):
                if df['close'].iloc[j + 1] > df['close'].iloc[j]:
                    up_days_count += 1

            psy_value = up_days_count / n * 100
            PSY[i] = psy_value

        return PSY

    def SAR(self, n=4, initial_af=0.02, max_af=0.2):
        """计算抛物线(SAR)指标"""
        self._validate_data(self.data, ['high', 'low'])
        df = self.data

        # 检查数据是否足够计算SAR值
        if len(df) <= n:
            return np.full(len(df), np.nan)

        # 提取高点和低点
        highs = df['high'].to_numpy()
        lows = df['low'].to_numpy()

        # 初始化存储SAR值和加速因子的数组
        sar_values = np.full_like(highs, np.nan, dtype=float)
        acceleration_factors = np.full_like(highs, initial_af, dtype=float)

        # 确定初始SAR值 (SAR(T0))
        if highs[n] > highs[0]:
            sar_values[n] = lows[0]
        else:
            sar_values[n] = highs[0]

        # 确定EP (极点价) 的初始值
        if highs[n] > highs[0]:
            ep = max(highs[1:n + 1])
        else:
            ep = min(lows[1:n + 1])

        # 计算每个周期的SAR值
        for i in range(n + 1, len(df)):
            if highs[i - n:i].max() > ep:
                ep = highs[i - n:i].max()
            elif lows[i - n:i].min() < ep:
                ep = lows[i - n:i].min()

            if highs[i - 1] > highs[i - n:i - 1].max() and sar_values[i - 1] <= highs[i - 1]:
                sar_values[i] = highs[i - 1]
                acceleration_factors[i] = initial_af
            elif lows[i - 1] < lows[i - n:i - 1].min() and sar_values[i - 1] >= lows[i - 1]:
                sar_values[i] = lows[i - 1]
                acceleration_factors[i] = initial_af
            else:
                sar_values[i] = sar_values[i - 1] + acceleration_factors[i] * (ep - sar_values[i - 1])
                acceleration_factors[i] = min(acceleration_factors[i - 1] + initial_af, max_af)

        return sar_values

    def OBOS(self, n=5):
        """计算OBOS涨跌指标"""
        self._validate_data(self.data, ['close'])
        df = self.data

        # 检查数据是否足够计算OBOS值
        if len(df) <= n:
            return np.full(len(df), np.nan)

        # 提取涨跌情况
        changes = df['close'].diff().dropna().to_numpy()

        # 初始化存储OBOS值的数组
        obos_values = np.full(len(df), np.nan)

        # 计算每日OBOS值
        for i in range(n, len(df)):
            # 统计N日内上涨和下跌家数
            na_sum = np.sum(changes[i - n + 1:i + 1] > 0)
            nb_sum = np.sum(changes[i - n + 1:i + 1] < 0)

            # 计算OBOS值
            obos_values[i] = na_sum - nb_sum

        return obos_values

    def CCI(self, n=20):
        """计算CCI指标"""
        self._validate_data(self.data, ['high', 'low', 'close'])
        df = self.data

        # 检查数据是否足够计算CCI值
        if len(df) <= n:
            return np.full(len(df), np.nan)

        # 计算典型价格（TP）
        tp = (df['high'] + df['low'] + df['close']) / 3

        # 计算MA（N日收盘价的简单移动平均）
        ma = tp.rolling(window=n, min_periods=1).mean()

        # 计算MD（N日的平均绝对偏差）
        md = np.abs(ma - df['close']).rolling(window=n, min_periods=1).mean()

        # 计算CCI值
        cci = []
        for i in range(len(md)):
            if md[i] == 0:
                cci.append(np.nan)
            else:
                cci.append((tp[i] - ma[i]) / (0.015 * md[i]))

        return cci

    def WVAD(self, n=24, m=6):
        "计算MWVAD"
        self._validate_data(self.data, ['open', 'high', 'low', 'close', 'vol'])

        df = self.data

        # 检查数据是否足够计算WVAD和MAWVAD
        if len(df) <= n:
            return np.full(len(df), np.nan)

        # 计算A和B
        A = df['close'] - df['open']
        B = df['high'] - df['low']

        # 计算C
        C_tmp = []
        for i in range(len(B)):
            if B[i] == 0:
                C_tmp.append(np.nan)
            else:
                C_tmp.append(A[i] / B[i])
        # C_tmp = np.where(B==0,np.nan,A/B)
        C = C_tmp * df['vol']

        # 计算WVAD
        WVAD = C.rolling(window=n, min_periods=1).sum()

        # 计算MAWVAD
        MAWVAD = WVAD.rolling(window=m, min_periods=1).mean()

        return MAWVAD

    def TAPI(self):
        """计算TAPI指标"""
        self._validate_data(self.data, ['amount'])
        df = self.data

        tapi_values = np.full(len(df), np.nan)

        # 计算每日OBOS值
        for i in range(1, len(df)):
            # 统计N日内上涨和下跌家数
            a = df['amount'].iloc[i]
            b = df['amount'].iloc[i - 1]

            if (b == 0):
                tapi_values[i] = 1e3

            tapi_values[i] = a + a / b * 100

        return tapi_values


if __name__ == '__main__':
    df = pd.read_csv('000008.SZ.csv')
    calculator = StockMetricsCalculator(df)
    # 测试SMA方法
    sma = calculator.SMA(window=5)
    print("简单移动平均线 (SMA):")
    print(sma)

    # 测试EMA方法
    ema = calculator.EMA(window=5)
    print("\n指数移动平均线 (EMA):")
    print(ema)

    # 测试RSI方法
    rsi = calculator.RSI(period=14)
    print("\n相对强弱指数 (RSI):")
    print(rsi)

    # 测试MACD方法
    macd_line, signal_line, histogram = calculator.MACD()
    print("\nMACD线:")
    print(macd_line)
    print("\n信号线:")
    print(signal_line)
    print("\n柱状图:")
    print(histogram)

    # 创建一个包含3行1列的子图布局
    fig, axes = plt.subplots(3, 1, figsize=(12, 10), sharex=True)

    # 绘制收盘价、SMA和EMA
    axes[0].plot(df['close'], label='Close Price')
    axes[0].plot(sma, label='SMA(5)')
    axes[0].plot(ema, label='EMA(5)')
    axes[0].set_title('Close Price, SMA(5) and EMA(5)')
    axes[0].set_ylabel('Price')
    axes[0].legend()

    # 绘制RSI
    axes[1].plot(rsi, label='RSI(14)')
    axes[1].axhline(y=30, color='r', linestyle='--', label='Oversold (30)')
    axes[1].axhline(y=70, color='g', linestyle='--', label='Overbought (70)')
    axes[1].set_title('Relative Strength Index (RSI)')
    axes[1].set_ylabel('RSI Value')
    axes[1].legend()

    # 绘制MACD指标
    axes[2].plot(macd_line, label='MACD Line')
    axes[2].plot(signal_line, label='Signal Line')
    axes[2].bar(histogram.index, histogram, label='Histogram', color='gray', alpha=0.5)
    axes[2].set_title('MACD Indicator')
    axes[2].set_xlabel('Date')
    axes[2].set_ylabel('MACD Value')
    axes[2].legend()

    plt.tight_layout()
    plt.show()